# TODO: At some point incorporate the useful bits of this directly in pycbg.

from random import choice
from random import shuffle
import shelve, sys, math, time
#
# These functions are all defined for ERC sets as lists of strings.
#
#==============================================================================#
#======================== Generate ERC sets ===================================#    

def randomERC(k):
    '''Generate a random nontrivial length-k ERC.'''
    R = [choice('lew') for x in range(k)]
    if 'l' in R and 'w' in R: return ''.join(R)
    else: return randomERC(k)

def randomERCs(k,n):
    '''Generate a list of n k-length nontrivial ERCs.'''
    return [randomERC(k) for x in range(n)]

#- These next two are good for profiling random consistent sets and
#- for exploring a permutation-representative slice of the ercSpace.

def randomConsistentEset(k,n):
    E = randomERCs(k,n)
    if consistent(E): return E
    else: return randomConsistentEset(k,n)

def randomERC_linear(k):
    '''Generate a random nontrivial length-k ERC in the linear order.'''
    R = [choice('lew') for x in range(k)]
    try: W,L = R.index('w'),R.index('l')
    except ValueError: W,L = 1,0
    if W < L: return ''.join(R)
    else: return randomERC(k)

def randomLinearERCs(k,n):
    '''Generate a list of n k-length nontrivial ERCs.'''
    return [randomERC_linear(k) for x in range(n)]

#==============================================================================#
#======================== Permute ERC sets ====================================#

def permPair(k,n):
    '''Returns two permutations of the same ERC set.'''
    X = randomERCs(k,n)
    Y = eshuf(X)
    return X,Y

def eshuf(elist):
    'Randomly shuffles the rows and columns of elist.'
    C = range(len(elist[0]))
    R = range(len(elist))
    shuffle(C); shuffle(R)
    return ePerm(elist,(C,R))

def ePerm(elist,perm):
    '''Permutes rows and columns according to perm.'''
    C,R = perm
    if not C or not R: return elist
    try: return [''.join([elist[r][c] for c in C]) for r in R]
    except IndexError:
        print 'CR:',C,R
        print elist
        return elist
        
#==============================================================================#
#===================== Basic ERC Functions ====================================#

def fuse(e1,e2):
    R = ''
    for (a,b) in zip(e1,e2):
        if a =='l' or b =='l': R+='l'
        elif a=='e' and b=='e': R+='e'
        else: R+= 'w'
    return R

def fusion(eset):
    '''Fusion of an entire ERC set.'''
    return reduce(fuse,eset)

def entails(e1,e2):
    '''Returns True if e1 -> e2.'''
    for (a,b) in zip(e1,e2):
        if b == 'l' and a != 'l': return False
        if b == 'e' and a == 'w': return False
    return True

def nontrivial(erc):
    return 'w' in erc and 'l' in erc

def consistent(eset):
    '''RDC-style consistency check.'''
    if len(eset) > 1: 
        fused = fusion(eset)
        if 'w' in fused:
            nu = [e for e in eset if not ('w','w') in zip(e,fused)]  
            return consistent(nu)
        else: return False
    else: return True

def strata(eset):
    '''Returns strata if eset is consistent [] otherwise.'''
    working = list(eset)
    strats = [set(range(len(eset[0])))] 
    while working:
        fused = reduce(fuse,working)  
        if 'w' in fused:
            layer = set([ix for (ix,X) in enumerate(fused) if X=='l'])
            strats[-1] -= layer
            strats += [layer]
            working = [e for e in working if not ('w','w') in zip(e,fused)]  
        elif 'l' in fused: return []
        else: break
    return [list(s) for s in strats]

def linearize(eset):
    '''R returns Eset that expresses a random total ordering consistent with eset.'''
    k = len(eset[0])
    S = strata(eset)
    for s in S: shuffle(s)
    S = oneflat(S)
    P = zip(S[:-1],S[1:])
    R = []
    for (a,b) in P:
        e = ''
        for x in range(k):
            if x == a: e += 'w'
            elif x == b: e+= 'l'
            else: e += 'e'
        R += [e]
    return R
            
#==============================================================================#
#==================== Useful Functions ========================================#

def llist(x):
    for q in x: print q

def oneflat(listlist):
    '''List of lists to list.'''
    return [x for sublist in listlist for x in sublist] 

def flatten(seq, a=[]):
    '''Nested iterables to list.'''
    if(isiterable(seq)):
        for item in seq: flatten(item, a)
    else: a.append(seq)
    return a

def sList2string(elist):
    '''Turns elist into a comma separated string.'''
    return ','.join(elist)

def unstringElist(estring):
    return estring.split(',')

def edis(eset):
    '''Display ERC matrices.'''
    for x in eset:
        for y in x: print y,
        print ''

def cross(listlist):
    ans = [[]]
    for L in listlist: ans = [x+[y] for x in ans for y in L]
    return ans

def allPerms(alist):
    '''Returns all permuations of alist.'''
    sz = len(alist)
    if sz <= 1:
        return [alist]
    return [p[:i]+[alist[0]]+p[i:] for i in xra7nge(sz) for p in perm(alist[1:])]

#calculates factorial and for n <=1 n! = 1
fac = lambda n: reduce(lambda a,b: a*(b+1), range(n), 1)

#==============================================================================#
#=============== Permutation Normal Form ======================================#

def pnorf(elist):
    '''Iteratively refine row and column partitions to PNF.'''
    if not elist: return [],[]
    Cols, Rows = [range(len(elist[0]))], [range(len(elist))]
    newC,newR = True, True
    while newC or newR:
        Lr,Lc,RR,CC = len(Rows),len(Cols),[],[]
        if newC: #- Refine the rows
            for Group in Rows:
                Gdic = {}
                for r in Group:
                    K = [(sum([elist[r][c]=='l' for c in cP]),
                          sum([elist[r][c]=='w' for c in cP])) for cP in Cols]
                    try: Gdic[tuple(K)].append(r)
                    except KeyError: Gdic[tuple(K)] = [r]
                RR += [Gdic[x] for x in sorted(Gdic.keys())]
        if newR: #- Refine the columns
            for Group in Cols:
                Gdic = {}
                for c in Group:
                    K = [(sum([elist[r][c]=='l' for r in rP]),
                          sum([elist[r][c]=='w' for r in rP])) for rP in Rows]
                    try: Gdic[tuple(K)].append(c)
                    except KeyError: Gdic[tuple(K)] = [c]
                CC += [Gdic[x] for x in sorted(Gdic.keys())]
        newC, newR = len(CC)>Lc, len(RR)>Lr
        if newC: Cols = CC
        if newR: Rows = RR
        if not newR and not newC and len(Rows)<len(elist):
            newR = True #- No more refinements - Break Symmetry
            for (ix,rs) in enumerate(Rows):
                if len(rs) > 1:
                    Rows = Rows[:ix]+[[rs[0]]]+[rs[1:]]+Rows[ix+1:];break
    return [c for P in Cols for c in P],[r for P in Rows for r in P]

        
def testPairs(k,n,itz): # Check itz pairs of equivalent k x n ERC sets
    inc = float(itz)/50
    print'|----------------------------------------------100%|\n|',
    while itz:
        if itz%inc < 1: sys.stdout.write('.'); sys.stdout.flush()
        itz -= 1
        A, B = permPair(k,n)
        Ap, Bp = pnorf(A), pnorf(B)
        nfA, nfB = ePerm(A,Ap), ePerm(B,Bp)
        if nfA!=nfB:  print '\n\n',A,'\n',nfA,'\n--------\n',B,'\n',nfB;break
    sys.stdout.write('|')

#==============================================================================#
#=================== Remove simple entailements ===============================#

def reducERCs(elist):
    """reduce() -> ERCSet

    The reduction of an ERCSet E is { e | e in E and there is no e' in E
    such that e' entails e other than e = e' }.
    """
    working, Result = list(elist), []
    while working:
        e = working.pop()
        for idx in range(len(working)-1,-2,-1):
            if idx < 0: Result.append(e)
            elif entails(e, working[idx]): working.pop(idx)
            elif entails(working[idx], e): break
    return Result

#==============================================================================#
#=================== Pnorf with no simple Ents ================================#

def sNorf(elist):
    reducedERCs = reducERCs(elist)
    permKey = pnorf(reducedERCs)
    return ePerm(reducedERCs,permKey),permKey

#==============================================================================#
#=================== ERC set minimization =====================================#

#def emin(elist):
#    """ Minimizes consistent erc sets represented as lists."""
#    knownErcs = dict.fromkeys(elist,1)
#    working = [0] + elist[:]
#    e1 = working.pop()
#    idx = len(working)-1
#    while e1:
#        if idx == -1:
#            working.insert(0,e1)
#            e1,idx = working.pop(), len(working)-1
#            continue
#        e2 = working[idx]
#        if e2 == 0: idx -=1; continue
#        e3 = fuse(e1,e2)
#        kkey = (entails(e3,e1),entails(e3,e2),entails(e1,e3),entails(e2,e3))
#        if kkey[1] and kkey[2]: del working[idx]
#        elif kkey[0] and kkey[3]: e1,idx = working.pop(),len(working) 
#        elif kkey[1] or kkey[2]: pass
#        else:
#            if not knownErcs.has_key(e3): knownErcs[e3]=1; working.append(e3)
#            if kkey[3] or not any(kkey[:2]): knownErcs[e3] = 0
#            if not kkey[1]: del working[idx]
#            if not kkey[0]: e1,idx = working.pop(),len(working)
#        idx -= 1
#    return [e for e in working if knownErcs[e]]


def transfuse(elist):
    '''Transitive closure of pair fusion in elist.'''
    knownERCs = set(elist)
    Nu,Old = list(elist),[]
    while Nu:
        nunu = []
        All = Old+Nu
        for (i,x) in enumerate(Nu):
            for y in All[:i+len(Old)]:
                z = fuse(x,y)
                if z not in knownERCs:
                    knownERCs.add(z)
                    nunu.append(z)
        Nu, Old = nunu, All
    return All

def simplemin(elist):
    TC = transfuse(elist)
    return reducERCs(TC)

#==============================================================================#
#=================== Minimized and normalized =================================#

def RNF(elist,withKey=False):
    if not consistent(elist):
        if withKey: return [],[]
        else: return []
    TC = transfuse(elist)
    red = reducERCs(TC)
    prm = pnorf(red)
    if withKey: return ePerm(red,prm),prm[0]
    else: return ePerm(red,prm)
    
#==============================================================================#
#=================== For counting rankings ====================================#

def bruteCheckah(elist,verbose=0):
    '''Brute force enumeration of orderings to check against elist.'''
    K, R = len(elist[0]), 0
    orders = allPerms(range(K))
    for ord in orders:
        if goodOrd(elist,ord):
            R += 1
            if verbose: print ord
    return R

def rankat(erc):
    '''Number of rankings consistent with ERC = (k!*w)/polar.'''
    k = fac(len(erc))
    w = sum([x=='w' for x in erc])
    p = sum([x!='e' for x in erc])
    return (k*w)/p

def fix(elist,ix):
    '''Removes any E that has 'w' in coordinate ix and removes ix alltogether.'''
    return [e[:ix]+e[ix+1:] for e in elist if e[ix] == 'e']

def ed_Fuse(e1,e2):
    '''E dominant fusion - where w+e = e.'''
    R = ''
    for (a,b) in zip(e1,e2):
        if a =='l' or b =='l': R += 'l'
        elif a == 'w' and b == 'w': R += 'w'
        else: R += 'e'
    return R

def nontrivial(erc):
    return 'l' in erc and 'w' in erc
        
def r(elist, snorfed=False, verbose=0):
    elist = [x for x in elist if nontrivial(x)]
    if verbose:
        print 'calculating r of', elist,
    if not consistent(elist):
        if verbose:
            print 0
        return (0,0)
    if len(elist) == 1: return rankat(elist[0]),1
    #Nonzero results for inconsistent ERC sets like:
    #['lewwe', 'lweee', 'leewl', 'weelw', 'eewll', 'lewee']
    # Should be fixed now.

    if not snorfed: elist = sNorf(elist)[0]
    todo, R, Fcount = {sList2string(elist) : [1,elist]}, 0, 0
    while todo:
        [n,E] = sorted(todo.values())[0]
        if verbose:
            print 'num',n
            edis(E)
        del todo[sList2string(E)]
        if len(E) == 1: R += (rankat(E[0])*n)
        else:
            F = reduce(ed_Fuse,E)
            Fcount += 1
            L = len(F)
            for ix in range(L):
                if F[ix] == 'w': R += n*(fac(L-1))
                elif F[ix] == 'e':
                    Nu = sNorf(fix(E,ix))[0]
                    sNu = sList2string(Nu)
                    try: todo[sNu][0] += n
                    except KeyError: todo[sNu] = [n,Nu]
    if verbose:
        print R
    return R, Fcount

def r_mem(elist):
    '''Memoized version of r computation'''
    rDict = shelve.open('rDict')
    snorfed = sNorf(elist)[0]
    E = sList2string(snorfed)
    if E in rDict:
        #sys.stderr.write('o'); sys.stderr.flush()
        R,Fcount = rDict[E]
    else:
        sys.stderr.write('+'); sys.stderr.flush()
        R, Fcount = r(snorfed,True)
        rDict[E] = R, Fcount
    rDict.close()
    return R, Fcount

#==============================================================================#
#=================== Profiling r ==============================================#

def prof_r(k,n,itz):
    '''Average and worst-case steps finding r in kXk matrices in itz trials.'''
    allsteps = 0
    print 'Doing',itz,'iterations of',k,'by',n,'ERC matrices:',
    while itz:
        #sys.stdout.write('.'); sys.stdout.flush()
        itz -= 1
        X = randomLinearERCs(k,n)
        rv,f = r_mem(X)
        allsteps += f
    return allsteps

def profile(itz):
    Profiles = shelve.open('r_profiles')
    if 'namedex' not in Profiles:
        Profiles['namedex']='0'
    Profiles.close()
    cdown = int(itz)
    while cdown>0:
        print cdown,'trials left in this session.'
        cdown -= 1
        for k in [10,20,30,40]:
            for x in range(1,26):
                Profiles = shelve.open('r_profiles')
                if k == 10: itx = 2
                if k == 20: itx = 10
                if k == 30: itx = 30
                if k == 40: itx = 60
                f = prof_r(k,x,itx) 
                print 'x'.join([str(k),str(x)]),f,itx
                Profiles[str(time.time())] = (k,x,f,itx)
                Profiles.close()
        r_report()

def r_report():
    import pylab
    rd = shelve.open('r_profiles')
    ten = ([x[1:] for x in rd.values() if x[0] == 10])
    twe = ([x[1:] for x in rd.values() if x[0] == 20])
    thr = ([x[1:] for x in rd.values() if x[0] == 30])
    fur = ([x[1:] for x in rd.values() if x[0] == 40])
    rd.close()
    end = 25
    v10 = [[0,0] for x in range(end)]
    for (x,f,i) in ten:
            if x <= end: v10[x-1][0] += f; v10[x-1][1] += i
    v10 = [a/float(b) for [a,b] in v10]

    v20 = [[0,0] for x in range(end)]
    for (x,f,i) in twe:
        if x <= end: v20[x-1][0] += f; v20[x-1][1] += i
    v20 = [a/float(b) for [a,b] in v20]

    v30 = [[0,0] for x in range(end)]
    for (x,f,i) in thr:
        if x <= end: v30[x-1][0] += f; v30[x-1][1] += i
    v30 = [a/float(b) for [a,b] in v30]

    v40 = [[0,0] for x in range(end)]
    for (x,f,i) in fur:
        if x <= end: v40[x-1][0] += f; v40[x-1][1] += i
    v40 = [a/float(b) for [a,b] in v40]

    pylab.clf()
    ax1 = pylab.subplot(111)
    pylab.plot([0]+v10)
    pylab.plot([0]+v20)
    pylab.plot([0]+v30)
    pylab.plot([0]+v40)
    #pylab.grid(True)
    ax1.yaxis.set_major_locator(pylab.MultipleLocator(5))
    pylab.text(4,4, '10 constraints', color='k')
    pylab.text(4,11,'20 constraints', color='k')
    pylab.text(4,24,'30 constraints', color='k')
    pylab.text(4,46,'40 constraints', color='k')
    pylab.xlabel('Number of ERCs')
    pylab.ylabel('Average number of recursive steps in computing r')
    pylab.show()




#==============================================================================#
#=================== generate test data =======================================#

def test_data(number):
    Efile = open('erctestdata.py','w+')
    Efile.write('testData = [')
    while number:
        number -=1
        N = choice(range(2,20))
        K = choice(range(3,30))
        E = randomERCs(K,N)
        Efile.write(repr(E))
        if number: Efile.write(',\n            ')
    Efile.write(']')
    Efile.close()

#==============================================================================#
#=================== ERC space with minimized normalized ERCs =================#

def monzy(k,n,itz):
    known = {}
    while itz:
        itz -= 1
        E = RNF(randomERCs(k,n))
        if E:
            S = sList2string(E)
            if S not in known: known[S] = len(E)     
    for E in [S for S in known if known[S]==max(known.values())]:
        E = unstringElist(E)
        edis(E)
        print '-'*(k+1),r(E)

#- RNF seems to be irreducable. For example,
#- these 2 denote different sets of rankings. 
s1 = ['ewel', 'weel', 'wwll'] 
s2 = ['ewel', 'wele', 'wwll']


#==============================================================================#
#=================== Extracting ERCs from ERC-set-disjunctions ================#

def ercOR(e1,e2):
    '''Returns the ERC that includes e1 or e2'''
    R = ''
    for (a,b) in zip(e1,e2):
        if a == 'w' or b == 'w' : R += 'w'
        elif a == 'l' and b == 'l' : R += 'l'
        else: R += 'e'
    return R
    
def ercSetOR(E1,E2):
    '''Returns the nontrivial ercs entialed by E1 or E2'''
    R = set([])
    for e1 in E1:
        for e2 in E2: 
            R.add(ercOR(e1,e2)) 
    return reducERCs([e for e in R if 'l' in e])


# e_1 = ['wwele', 'eweel', 'ewwll'] 
# e_2 = ['ewwel', 'ewele', 'ewwll']
# 
# print ercSetOR(e_1,e_2)